# Copyright (c) 2025, Huawei Technologies Co., Ltd. All rights reserved.
import os
import random

import numpy as np
import torch


def set_seed(seed: int, set_deterministic=False):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.accelerator.manual_seed(seed)
    torch.accelerator.manual_seed_all(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    if set_deterministic:
        set_deterministic_algorithms()


def set_deterministic_algorithms():
    '''
    HCCL_DETERMINISTIC: is a deterministic switch in ops level, set it to 'True' to enable ops level deterministic,
        set it to 'False' to disable ops level deterministic.
    CLOSE_MATMUL_K_SHIFT: is a switch of matmul K-axis shift, set it to '1' to close matmul K-axis shift,
        set it to '0' to enable matmul K-axis shift.
    '''
    os.environ['HCCL_DETERMINISTIC'] = 'True'
    os.environ['CLOSE_MATMUL_K_SHIFT'] = '1'
    torch.use_deterministic_algorithms(True)
